import sys
from PyQt5.QtWidgets import (QApplication, QWidget, QVBoxLayout, QHBoxLayout,
                             QLabel, QLineEdit, QComboBox, QPushButton,
                             QTextEdit, QStatusBar)
from PyQt5.QtCore import QThread, pyqtSignal, Qt, QTimer
from PyQt5.QtNetwork import QTcpSocket, QUdpSocket, QHostAddress
from PyQt5.QtGui import QTextCursor


class NetworkDebugTool(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("网络调试助手")
        self.setGeometry(100, 100, 800, 600)

        # 网络组件
        self.tcp_socket = QTcpSocket()
        self.udp_socket = QUdpSocket()

        self.init_ui()
        self.init_connections()

    def init_ui(self):
        # 主布局
        main_layout = QVBoxLayout()

        # 协议选择
        protocol_layout = QHBoxLayout()
        protocol_layout.addWidget(QLabel("协议:"))
        self.protocol_combo = QComboBox()
        self.protocol_combo.addItems(["TCP", "UDP"])
        protocol_layout.addWidget(self.protocol_combo)

        # 地址和端口输入
        addr_layout = QHBoxLayout()
        addr_layout.addWidget(QLabel("目标地址:"))
        self.addr_input = QLineEdit("127.0.0.1")
        addr_layout.addWidget(self.addr_input)

        addr_layout.addWidget(QLabel("端口:"))
        self.port_input = QLineEdit("8080")
        addr_layout.addWidget(self.port_input)

        # 连接控制按钮
        btn_layout = QHBoxLayout()
        self.connect_btn = QPushButton("连接")
        self.disconnect_btn = QPushButton("断开")
        self.disconnect_btn.setEnabled(False)
        btn_layout.addWidget(self.connect_btn)
        btn_layout.addWidget(self.disconnect_btn)

        # 状态指示灯
        self.status_light = QLabel("◌")
        self.status_light.setStyleSheet("font-size: 24px;")
        btn_layout.addWidget(self.status_light)

        # 数据发送区域
        send_layout = QVBoxLayout()
        send_layout.addWidget(QLabel("发送数据:"))
        self.send_text = QTextEdit()
        send_layout.addWidget(self.send_text)

        send_btn_layout = QHBoxLayout()
        self.send_btn = QPushButton("发送")
        self.send_btn.setEnabled(False)
        self.clear_send_btn = QPushButton("清空")
        send_btn_layout.addWidget(self.send_btn)
        send_btn_layout.addWidget(self.clear_send_btn)
        send_layout.addLayout(send_btn_layout)

        # 数据接收区域
        recv_layout = QVBoxLayout()
        recv_layout.addWidget(QLabel("接收数据:"))
        self.recv_text = QTextEdit()
        self.recv_text.setReadOnly(True)
        recv_layout.addWidget(self.recv_text)

        recv_btn_layout = QHBoxLayout()
        self.clear_recv_btn = QPushButton("清空")
        self.save_recv_btn = QPushButton("保存")
        recv_btn_layout.addWidget(self.clear_recv_btn)
        recv_btn_layout.addWidget(self.save_recv_btn)
        recv_layout.addLayout(recv_btn_layout)

        # 状态栏
        self.status_bar = QStatusBar()
        self.status_bar.showMessage("准备就绪")

        # 组合所有布局
        main_layout.addLayout(protocol_layout)
        main_layout.addLayout(addr_layout)
        main_layout.addLayout(btn_layout)
        main_layout.addLayout(send_layout)
        main_layout.addLayout(recv_layout)
        main_layout.addWidget(self.status_bar)

        self.setLayout(main_layout)

    def init_connections(self):
        # 按钮信号连接
        self.connect_btn.clicked.connect(self.connect_to_host)
        self.disconnect_btn.clicked.connect(self.disconnect_from_host)
        self.send_btn.clicked.connect(self.send_data)
        self.clear_send_btn.clicked.connect(lambda: self.send_text.clear())
        self.clear_recv_btn.clicked.connect(lambda: self.recv_text.clear())
        self.save_recv_btn.clicked.connect(self.save_received_data)

        # TCP信号连接
        self.tcp_socket.connected.connect(self.on_connected)
        self.tcp_socket.disconnected.connect(self.on_disconnected)
        self.tcp_socket.errorOccurred.connect(self.on_error)
        self.tcp_socket.readyRead.connect(self.tcp_receive_data)

        # UDP信号连接
        self.udp_socket.readyRead.connect(self.udp_receive_data)

    def connect_to_host(self):
        host = self.addr_input.text()
        port = int(self.port_input.text())
        protocol = self.protocol_combo.currentText()

        if protocol == "TCP":
            self.tcp_socket.abort()
            self.tcp_socket.connectToHost(host, port)

            # 设置连接超时
            QTimer.singleShot(5000, lambda:
            self.update_status(False) if self.tcp_socket.state() == QTcpSocket.ConnectingState else None)
        else:  # UDP
            self.update_status(True)
            self.status_bar.showMessage(f"UDP已准备好发送数据到 {host}:{port}")

        self.connect_btn.setEnabled(False)
        self.disconnect_btn.setEnabled(True)
        self.send_btn.setEnabled(True)

    def disconnect_from_host(self):
        protocol = self.protocol_combo.currentText()

        if protocol == "TCP":
            self.tcp_socket.abort()
        else:  # UDP
            pass  # UDP无需断开连接

        self.on_disconnected()

    def on_connected(self):
        self.update_status(True)
        self.status_bar.showMessage(f"已连接到 {self.tcp_socket.peerName()}:{self.tcp_socket.peerPort()}")

    def on_disconnected(self):
        self.update_status(False)
        self.status_bar.showMessage("已断开连接")
        self.connect_btn.setEnabled(True)
        self.disconnect_btn.setEnabled(False)
        self.send_btn.setEnabled(False)

    def on_error(self, error):
        self.status_bar.showMessage(f"错误: {self.tcp_socket.errorString()}")
        self.update_status(False)

    def update_status(self, is_connected):
        color = "green" if is_connected else "red"
        text = "●" if is_connected else "◌"
        self.status_light.setStyleSheet(f"color: {color}; font-size: 24px;")
        self.status_light.setText(text)

    def send_data(self):
        data = self.send_text.toPlainText()
        if not data:
            return

        host = self.addr_input.text()
        port = int(self.port_input.text())
        protocol = self.protocol_combo.currentText()

        if protocol == "TCP":
            if self.tcp_socket.state() == QTcpSocket.ConnectedState:
                self.tcp_socket.write(data.encode())
        else:  # UDP
            self.udp_socket.writeDatagram(data.encode(), QHostAddress(host), port)

        # 在接收区显示发送的数据
        self.append_to_recv(f"[发送] {data}", color="blue")

    def tcp_receive_data(self):
        while self.tcp_socket.bytesAvailable():
            data = self.tcp_socket.readAll().data().decode()
            self.append_to_recv(f"[接收] {data}", color="green")

    def udp_receive_data(self):
        while self.udp_socket.hasPendingDatagrams():
            datagram, host, port = self.udp_socket.readDatagram(
                self.udp_socket.pendingDatagramSize())
            data = datagram.decode()
            self.append_to_recv(f"[接收] {data} (来自 {host}:{port})", color="green")

    def append_to_recv(self, text, color=None):
        cursor = self.recv_text.textCursor()
        cursor.movePosition(QTextCursor.End)

        if color:
            self.recv_text.setTextColor(Qt.black)
            self.recv_text.insertPlainText("\n")
            self.recv_text.setTextColor(Qt.darkGreen if color == "green" else Qt.blue)

        self.recv_text.insertPlainText(text)
        self.recv_text.ensureCursorVisible()

    def save_received_data(self):
        # 这里可以添加保存接收数据到文件的逻辑
        pass

    def closeEvent(self, event):
        self.tcp_socket.abort()
        self.udp_socket.close()
        event.accept()


if __name__ == "__main__":
    app = QApplication(sys.argv)
    tool = NetworkDebugTool()
    tool.show()
    sys.exit(app.exec_())